import json as json
import numpy
import numpy as np
from matplotlib.ticker import MaxNLocator
from matplotlib import cm
from PIL import Image
import io


def genTick(imin, imax, splitN, prune=None):
  precise = 3
  loc = MaxNLocator(splitN, prune=prune)
  imin = round(imin, precise)
  imax = round(imax, precise)
  ticks = loc.tick_values(imin, imax)
  ticks = numpy.around(ticks, decimals = 3)
  return dict(min=imin, max=imax, ticks=ticks)


FREQ_META = dict(
  l=genTick(130, 170, 8),
  h=genTick(300, 1500, 8)
)
def calc_sample(data, normal=True):
  if len(data) < 1: return []
  tmp_data = numpy.abs(data).astype('f8')
  if normal:
    tmp_data = tmp_data/numpy.nanmax(tmp_data)
  return 10*numpy.log10(tmp_data ** 2)


dtype = [('lxm', 'f4'), ('pul', 'c8'), ('amp', 'f4'), ('pha', 'f4'), ('real', 'f4'), ('sample', 'c8')]

# 历史绘图
async def hist(request, raw, **kwargs):
  print('-------------------------------img begin', flush=True)
  t = request.args.t
  mode = int(request.args.mode)
  key = request.args.k

  cf = []
  with open(f'{t}/.info', 'r') as f:
    for i in f:
      cf.append(json.loads(i))

  cf = [i for i in cf[1:] if i['mode'] == mode]
  o = []
  # N = 10
  N = int(len(cf) / 600.) or 1
  with open(f'{t}/.dat', 'rb') as f:
    tmp = []
    for i in cf:
      off, L, H = i['frm']
      f.seek(off, 0)
      t = dict(
        l_h=numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201),
        l_v=numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201),
        h_a=numpy.fromfile(f, dtype, count=H * 6001).reshape(H, 6001),
      )
      tmp.append(t[key]['sample'])
      if len(tmp) == N:
        tmp = numpy.concatenate(tmp, axis=0)
        shp = tmp.shape
        tmp = tmp.mean(axis=0)
        tmp = numpy.log10(numpy.abs(tmp) ** 2)
        o.append(tmp)
        tmp = []
    if len(tmp):
      tmp = numpy.concatenate(tmp, axis=0)
      shp = tmp.shape
      tmp = tmp.mean(axis=0)
      tmp = numpy.log10(numpy.abs(tmp) ** 2)
      o.append(tmp)
  o = np.array(o).T
  o[o==np.inf] = np.nan
  o[o==-np.inf] = np.nan
  print(f'-------o.shape: {o.shape}')
  fram = genTick(1, len(cf), 8, prune='upper')
  sample = [1,201] if key[0]=='l' else [1,6001]
  sample_ticks = genTick(sample[0], sample[-1], (4 if key[0]=='l' else 6), prune='upper')
  sample_ticks['ticks'][0] = 1
  a = numpy.nanmin(o)
  b = numpy.nanmax(o)
  cbar = genTick(a, b, 8, prune='both')

  cmp = cm.get_cmap(name='jet')
  lut = cmp(range(256), bytes=True)[:, :3]

  o -= a
  o *= 255 / (b - a)
  r = lut.take(o.astype('i4'), mode='clip', axis=0)

  with io.BytesIO() as f:
    Image.fromarray(r).save(f, format='png')
    v = f.getvalue()
  info = json.dumps(dict(N=N, sample=sample_ticks, fram=fram, cbar=cbar), default=tolist)
  return raw(v, headers={'content-type': 'image/png', 'info': info})


def tolist(v):
  if hasattr(v, 'tolist'): return v.tolist()
  return v


def _c(d):
  # s_src = d['sample'].astype('f8')
  # s_src[numpy.isinf(s_src)] = numpy.nan

  s = calc_sample(d['sample'],True)
  a_src = d['amp'].astype('f8')
  p_src = d['pha'].astype('f8')
  r_src = d['real'].astype('f8')
  a_src[np.isinf(a_src)] = np.nan
  p_src[np.isinf(p_src)] = np.nan
  r_src[np.isinf(r_src)] = np.nan
  return dict(
    amp=numpy.around(a_src, decimals=2).tolist(),
    pha=numpy.around(p_src, decimals=2).tolist(),
    real=numpy.around(r_src, decimals=2).tolist(),
    sample=numpy.around(s, decimals=2).tolist()
  )


async def rt(request, json, **kwargs):
  fp = request.args.f
  offset = int(request.args.offset)
  L = int(request.args.L)
  H = int(request.args.H)
  dtype = [('lxm', 'f4'), ('pul', 'c8'), ('amp', 'f4'), ('pha', 'f4'), ('real', 'f4'), ('sample', 'c8')]
  with open(f'{fp}/.dat', 'rb') as f:
    f.seek(offset, 0)
    L_H = numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201)
    L_V = numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201)
    H_A = numpy.fromfile(f, dtype, count=H * 6001).reshape(H, 6001)

  return json(dict(
    l_h=_c(L_H),
    l_v=_c(L_V),
    h_a=_c(H_A),
  ))

def _acc_c(d_arr):
  d = numpy.concatenate(d_arr, axis=0)
  l = len(d)
  print(f'_acc_c d.shape:{d.shape}')
  if l:
    rst = dict(v0=numpy.around(calc_sample(d[0], False),decimals=2).tolist())
    if l > 1:
      d_mean = d.mean(axis=0)
      vn = numpy.around(calc_sample(d_mean, False), decimals=2).tolist()
      rst.update(dict(vn=vn, n=d.shape[0]))
      # vn = numpy.around(numpy.sum(d, axis=0), decimals=2).tolist()
      # rst.update(dict(vn=vn, n=d.shape[0]))
    return rst
  return []

# 历史非噪声模式接口
async def hist_chart(request, json, **kwargs):
  params = request.json
  infos = params['infos']
  fp = params['f']

  LH_list = []
  LV_list = []
  HA_list = []
  with open(f'{fp}/.dat', 'rb') as f:
    for info in infos:
      offset, L,H = info['frm']
      offset = int(offset)
      L = int(L)
      H = int(H)
      f.seek(offset, 0)
      LH_list.append(numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201)['sample'])
      LV_list.append(numpy.fromfile(f, dtype, count=L * 201).reshape(L, 201)['sample'])
      HA_list.append(numpy.fromfile(f, dtype, count=H * 6001).reshape(H, 6001)['sample'])

  return json(dict(
    l_h=_acc_c(LH_list),
    l_v=_acc_c(LV_list),
    h_a=_acc_c(HA_list),
  ))

# 历史噪声模式
import json as tjson
from pathlib import Path
# 历史噪声模式
async def hist_voice(request, json, **kwargs):
  fp = request.args.f
  mode = int(request.args.m)
  N = 100
  def acc_mean(d):
    tmp = numpy.concatenate(d, axis=0)
    if len(tmp):
      return tmp.mean(axis=0)
    return numpy.array([])
  info_fp = Path(f'{fp}/.info')
  data_fp = Path(f'{fp}/.dat')
  L_LEN = 201
  H_LEN = 6001
  dtype = [('lxm', 'f4'), ('pul', 'c8'), ('amp', 'f4'), ('pha', 'f4'), ('real', 'f4'), ('sample', 'c8')]
  L_V = []
  L_H = []
  H_A = []
  if info_fp.exists() and data_fp.exists():
    with open(info_fp, 'r', encoding='UTF8') as f:
      with open(data_fp, mode='rb') as dataf:
        for line in f:
          try:
            info = tjson.loads(line)
          except:
            print(line)
          if 'mode' in info:
            imode = info['mode']
            if imode == mode:
              offset, L, H = info['frm']
              dataf.seek(offset)
              lv = numpy.fromfile(dataf, dtype=dtype, count=L_LEN * L).reshape(L, L_LEN)['lxm'].astype('f8')
              L_V.append(lv)
              L_H.append(numpy.fromfile(dataf, dtype=dtype, count=L_LEN * L).reshape(L, L_LEN)['lxm'].astype('f8'))
              H_A.append(numpy.fromfile(dataf, dtype=dtype, count=H_LEN * H).reshape(H, H_LEN)['lxm'].astype('f8'))
              if (len(L_V)) == N:
                L_V = [acc_mean(L_V).reshape(-1, L_LEN)]
                L_H = [acc_mean(L_H).reshape(-1, L_LEN)]
                H_A = [acc_mean(H_A).reshape(-1, H_LEN)]

      if len(L_V):
        L_V = numpy.around(10*numpy.log10(acc_mean(L_V)), decimals=3)
        L_H = numpy.around(10*numpy.log10(acc_mean(L_H)), decimals=3)
        H_A = numpy.around(10*numpy.log10(acc_mean(H_A)), decimals=3)
        L_V[numpy.isnan(L_V)] = -9999
        L_H[numpy.isnan(L_H)] = -9999
        H_A[numpy.isnan(H_A)] = -9999
  return json(dict(l_v=L_V.tolist(), l_h=L_H.tolist(), h_a=H_A.tolist()))